from typing import List, Any
import math
from ortools.linear_solver import pywraplp as lp


class GroupSolver:
    IDLE = 0
    OPTIMAL = 1
    FEASIBLE = 2
    INFEASIBLE = 3

    def __init__(self,
            area_staffs,    # 员工分区域数据 
            flights,        # 航班数据
            group_min=4,    # 组内最少员工数
            group_max=7,    # 组内最多员工数
            group_cnt=9,    # 推荐组数
            deviation=3,    # 组数偏移量
            ):
        self._area_staffs = area_staffs
        self._flights = flights
        self._group_min = group_min
        self._group_max = group_max
        self._group_cnt = group_cnt
        self._deviation = deviation

        self._res = None
        self._status = GroupSolver.IDLE
        self._warning = None

    def create_solver(self):
        # 初始化求解器
        solver = lp.Solver("GroupSolver", lp.Solver.CBC_MIXED_INTEGER_PROGRAMMING)
        
        # 变量
        var_data = {}
        for areaId, staff_id_list in self._area_staffs.items():
            # 最多分组数量
            max_group_cnt = math.ceil(len(staff_id_list)/self._group_min)
            var_data[areaId] = [{
                "decision_var": solver.IntVar(0, self._group_max, f"Area({areaId}) {i}-th Group"),
                "support_var": solver.IntVar(0, 1, f"Sup Area({areaId}) {i}-th Group")
                }
                for i in range(max_group_cnt)
            ]

        min_group_cnt_var =  solver.IntVar(0, self._group_max, "min_cnt")
        max_group_cnt_var =  solver.IntVar(0, self._group_max, "max_cnt")

        # 约束
        # 组员人数约束
        for areaId, group_var in var_data.items():
            for i in range(len(group_var)):
                solver.Add(group_var[i]["support_var"]*self._group_min 
                            <= group_var[i]["decision_var"])
                solver.Add(group_var[i]["decision_var"] <= 
                            group_var[i]["support_var"]*self._group_max)     

        # 区域总人数约束
        for areaId, group_var in var_data.items():
            solver.Add(solver.Sum([group_var[i]["decision_var"] for i in range(len(group_var))]) 
                == len(self._area_staffs[areaId]))

        # 总组数约束
        cons = solver.Constraint(self._group_cnt - self._deviation, self._group_cnt + self._deviation)
        for areaId, group_var in var_data.items():
            for i in range(len(group_var)):
                cons.SetCoefficient(group_var[i]["support_var"], 1)

        # 组最小最大人数变量的约束
        for areaId, group_var in var_data.items():
            for i in range(len(group_var)):
                solver.Add(min_group_cnt_var <= group_var[i]["decision_var"])
                solver.Add(max_group_cnt_var >= group_var[i]["decision_var"])

        # 目标
        obj = solver.Objective()
            # 组员人数平均
        obj.SetCoefficient(max_group_cnt_var, 1)
        obj.SetCoefficient(min_group_cnt_var, -1)
        #     # 组总数较少
        # for areaId, group_var in var_data.items():
        #     for i in range(len(group_var)):
        #         obj.SetCoefficient(group_var[i]["support_var"], 1)
        obj.SetMinimization()
        return solver, var_data

    def solve(self):
        while True:
            solver, var_data = self.create_solver()

            # 求解
            solver.EnableOutput()
            status = solver.Solve()
            if status == lp.Solver.OPTIMAL:
                self._status = GroupSolver.OPTIMAL

                # 计算结果数据
                staff_group_list = []
                group_id = 0
                for areaId, group_var in var_data.items():
                    cur_index = 0
                    staff_id_list = self._area_staffs[areaId]
                    for i in range(len(group_var)):

                        # print("====")
                        # print(f"{group_var[i]['support_var'].solution_value()*self._group_min} <= {group_var[i]['decision_var'].solution_value()} <= {group_var[i]['support_var'].solution_value()*self._group_max}")

                        if group_var[i]["support_var"].solution_value() == 1:
                            cur_gp_cnt = int(group_var[i]["decision_var"].solution_value())
                            for index in range(cur_index, cur_index + cur_gp_cnt):
                                staff_group_list.append((staff_id_list[index], group_id, areaId))
                            group_id += 1
                            cur_index += cur_gp_cnt
                self._res = staff_group_list

                print(f"Total Groups: {group_id}")
                return self._status
            else:
                self._group_min -= 1
                self._warning = f"组最小人数变更为 {self._group_min} 人"
                if not self._group_min:
                    break

        self._status = GroupSolver.INFEASIBLE
        return self._status

    @property
    def status(self):
        return self._status

    def calc_return_data(self):
        return self._res